from MCTS import MCTS
import time
import utils
from network import neuralnetwork as nn
import os
import torch


class Trainer:
    def __init__(self, board_size: int = 11, lr=0.001, pretrained_model=None, use_cuda: bool = True, game_file_saved_dict="game_record"):
        if not os.path.exists(game_file_saved_dict):
            os.mkdir(game_file_saved_dict)

        self.Net = torch.load(pretrained_model) if pretrained_model else nn(input_layers=3, board_size=board_size, learning_rate=lr, use_cuda=use_cuda)
        self.Net.device = "cuda" if use_cuda else "cpu"
        self.game_file_saved_path = game_file_saved_dict

    def train(self):
        stack = utils.random_stack()
        tree = MCTS(board_size=utils.board_size, neural_network=self.Net)
        record = []
        game_time = 0

        while True:
            game_record, eval, steps = tree.game()

            if len(game_record) % 2 == 1:
                print(f"game {game_time} completed, black win, this game length is {len(game_record)}")
            else:
                print(f"game {game_time} completed, white win, this game length is {len(game_record)}")
            print(f"The average eval:{eval}, the average steps:{steps}")

            utils.write_file(game_record, self.game_file_saved_path + "/" + time.strftime("%Y%m%d_%H_%M_%S", time.localtime()) + f"_{game_time}.pkl")
            train_data = utils.generate_training_data(game_record=game_record, board_size=utils.board_size)
            for i in range(len(train_data)):
                stack.push(train_data[i])
            loaded_data = utils.generate_data_loader(stack)
            utils.write_file(loaded_data, self.game_file_saved_path + "/debug_loader.pkl")
            if game_time % 100 == 0:
                for _ in range(5):
                    record.extend(self.Net.train(loaded_data, game_time))
            print("train finished")

            if game_time % 200 == 0:
                torch.save(self.Net, f"model/model_{game_time}.pkl")
                test_game_record, _, _ = tree.game(train=False)
                utils.write_file(test_game_record, self.game_file_saved_path + "/" + f"test_{game_time}.pkl")
                print(f"We finished a test game at {game_time} game time")
            game_time += 1


if __name__ == "__main__":
    trainer = Trainer(board_size=utils.board_size, pretrained_model="model_5400.pkl", use_cuda=True)
    trainer.train()
